#include "MBImageManager.h"
#include "MBDirectoryMethods.h"
#include "MBStringMethods.h"
#include "MBTimer.h"

#include <iostream>
#include <math.h>

using cv::Mat;
using cv::Size;

using namespace MB;

using std::cout;
using std::endl;
using std::make_pair;
using std::map;
using std::max;
using std::pair;
using std::string;
using std::vector;

MBImageManager *MBImageManager::imageManager= 0;

/**
  *	Constructs a new MBImageManager object with a maximum storage consumption
  * of 1GB in memory
  */
MBImageManager::MBImageManager( void )
{
	mCurrentImageOrder = 0u;
	mCurrentMemoryConsumption = 0u;
	mExclusiveLevel = 0u;

	mNumImagesLoaded = 0u;
	mSecondsNeededForLoading = 0.0;

	// Set maximum memory consumption to 1 GB = 1 * 1024 * 1024 * 1024 B = 1073741824 B
	SetMaximumMemoryConsumptionInBytes( 1073741824u );
}

/**
  *	Destroys the MBImageManager and all the images in memory
  */
MBImageManager::~MBImageManager( void )
{
	FreeAllImages();
}

/**
  *	Frees all the images so that the memory consumption is 0 afterwards
  */
void MBImageManager::FreeAllImages( void )
{
	for (map<string, vector<Mat*>>::iterator it = mImagePathToImagePyramidMap.begin();
		it != mImagePathToImagePyramidMap.end(); it++)
	{
		for (vector<Mat*>::iterator mat_it = it->second.begin();
			mat_it != it->second.end(); mat_it++)
		{
			delete *mat_it;
		}
	}

	mImagePathToImagePyramidMap.clear();
	mImagePathToCurrentOrderMap.clear();
	mCurrentOrderToImagePathMap.clear();
	mCurrentMemoryConsumption = 0;
}

/**
  *	Frees the image that belongs to the given image path and its entire image
  * pyramid
  *
  * @param imagePath path of an image file
  * @return number of bytes that could be freed
  */
MBuint MBImageManager::FreeImage( const std::string &imagePath )
{
	// cout << "Free Image: " << GetFileNameFromPath( imagePath ) << endl;
	map<string, vector<Mat*>>::iterator it = mImagePathToImagePyramidMap.find( imagePath );

	if (it != mImagePathToImagePyramidMap.end())
	{
		for (vector<Mat*>::iterator mat_it = it->second.begin();
			mat_it != it->second.end(); mat_it++)
		{
			if (*mat_it) delete *mat_it;
		}

		map<string, MBuint>::iterator im_to_order_it = mImagePathToCurrentOrderMap.find( imagePath );
		map<MBuint, MBImagePathWithMemoryConsumption>::iterator order_to_im_it = mCurrentOrderToImagePathMap.find( im_to_order_it->second );

		MBuint consumption = order_to_im_it->second.mMemoryConsumption;
		mCurrentMemoryConsumption -= consumption;

		mImagePathToImagePyramidMap.erase( it );
		mImagePathToCurrentOrderMap.erase( im_to_order_it );
		mCurrentOrderToImagePathMap.erase( order_to_im_it );

		return consumption;
	}

	return 0;
}

/**
  *	Returns the only instance of the MBImageManager class
  *
  * @return only instance of the MBImageManager class
  */
MBImageManager &MBImageManager::Get( void )
{
	if (!imageManager)
		imageManager = new MBImageManager;

	return *imageManager;
}

/**
  *	Returns a const pointer to the cv::Mat matrix that contains the image data
  * of the image, that belongs to the given image path and layer number
  *
  * @param imagePath path of an image
  * @param layerNumber number of the desired layer in the image pyramid, 0 is the original image
  * @param untouchableImagePaths set with all paths of images, that must not be freed when loading the desired image
  * @return const pointer to the cv::Mat matrix that contains the image data of the image, that belongs to the given image path and layer number
  */
const Mat *MBImageManager::GetImage(
	const string &imagePath,
	const MBuint layerNumber,
	const std::set<std::string> *untouchableImagePaths )
{
	MBTimer timer;

	Mat *image = 0;
	pair<map<string, vector<Mat*>>::iterator, bool> insert_pair;
	
	insert_pair = mImagePathToImagePyramidMap.insert( make_pair( imagePath, vector<Mat*>() ) );

	if (insert_pair.second)
	{
		// Insertion was successful, this is a new image -> load
		image = new Mat();
		*image = cv::imread( imagePath );

		if (!image->data)
		{
			// Invalid image path or not enough memory
			cout << "ERROR: Could not load " << GetFileNameFromPath( imagePath ) << endl;
			mImagePathToImagePyramidMap.erase( insert_pair.first );
			delete image;
			return 0;
		}
		else
		{	// Image successfully read
			mNumImagesLoaded++;

			//Mat *double_image = new Mat();
			MBuint image_memory_consumption = 0u;
			//image->convertTo( *double_image, CV_64FC3 );
			//delete image;

			// cout << "Load Image: " << GetFileNameFromPath( imagePath ) << " at level " << layerNumber << endl;

			// Bring the values into the [0 1] range
			//*double_image /= 255.0;

			if (layerNumber == 0 || mExclusiveLevel == 0)
			{
				insert_pair.first->second.push_back( image );
				image_memory_consumption += 3 * image->cols * image->rows * sizeof( unsigned char );
// 				insert_pair.first->second.push_back( double_image );
// 				image_memory_consumption += 3 * double_image->cols * double_image->rows * sizeof( double );
			}
			else
			{
				// Save Null-Pointer for Level 0
				insert_pair.first->second.push_back( 0 );
			}

			// Build the pyramid
			//Mat *src = double_image;
			Mat *src = image;
			Mat *dst;
			MBint current_level = 0;
			MBint last_layer = max( mExclusiveLevel, layerNumber );

			MBint exclusive_border = mExclusiveLevel - 1;
			MBint layer_border = layerNumber - 1;

			bool delete_source_image = false;

			while (true)
			{
				if (current_level < last_layer && src->cols >= 2 && src->rows >= 2)
				{
					dst = new Mat();
					pyrDown( *src, *dst, Size( src->cols / 2, src->rows / 2 ) );

					if (insert_pair.first->second.back() == 0)
					{
						// The last level was not needed -> delete source image
						delete_source_image = true;
					}

					if (current_level == exclusive_border || current_level == layer_border)
					{
						insert_pair.first->second.push_back( dst );
						image_memory_consumption += 3 * dst->cols * dst->rows * sizeof( unsigned char );
						//image_memory_consumption += 3 * dst->cols * dst->rows * sizeof( double );
					}
					else if (current_level < exclusive_border || current_level < layer_border)
					{
						// Save Null-Pointer for not needed Levels
						insert_pair.first->second.push_back( 0 );
					}

					if (delete_source_image)
					{
						delete src;
						delete_source_image = false;
					}

					src = dst;

					current_level++;
				}
				else
				{
					break;
				}
			}

			mImagePathToCurrentOrderMap.insert( make_pair( imagePath, mCurrentImageOrder ) );
			mCurrentOrderToImagePathMap.insert( make_pair( mCurrentImageOrder++, MBImagePathWithMemoryConsumption( imagePath, image_memory_consumption ) ) );

			// Check memory consumption, and free some if needed
			mCurrentMemoryConsumption += image_memory_consumption;

			MBuint tries = 0u;
			map<MBuint, MBImagePathWithMemoryConsumption>::iterator order_to_image_it = mCurrentOrderToImagePathMap.begin();

			while (mCurrentMemoryConsumption > mMaxMemoryConsumption && mImagePathToCurrentOrderMap.size() - tries > 1u)
			{
				if (untouchableImagePaths)
				{
					// There are images, that MUST NOT be deleted
					if (untouchableImagePaths->find( order_to_image_it ->second.mImagePath ) == untouchableImagePaths->end() )
					{
						// OK, this image can be deleted
						FreeImage( order_to_image_it->second.mImagePath );
						order_to_image_it = mCurrentOrderToImagePathMap.begin();
					}
					else
					{
						// This image must not be deleted, choose the next one
						order_to_image_it++;
					}
				}
				else
				{
					// No untouchable images, just delete
					FreeImage( mCurrentOrderToImagePathMap.begin()->second.mImagePath );
				}
				tries++;
			}

			if (layerNumber >= insert_pair.first->second.size())
			{
				return 0;
			}
			else
			{
				mSecondsNeededForLoading += timer.Toc();
				return insert_pair.first->second[layerNumber];
			}
		}
	}
	else
	{
		// Existing image

		if (layerNumber >= insert_pair.first->second.size())
		{
			if (insert_pair.first->second.back()->cols < 2u || insert_pair.first->second.back()->rows < 2u)
			{
				// We are on the highest mipmap level
				return 0;
			}
			else
			{
				MBint current_level = insert_pair.first->second.size() - 1;
				MBuint image_memory_consumption = 0u;
				
				// Build the needed layers of the pyramid
				Mat *src = insert_pair.first->second.back();
				Mat *dst;

				MBint last_layer = max( mExclusiveLevel, layerNumber );

				MBint exclusive_border = static_cast<MBint>( mExclusiveLevel ) - 1;
				MBint layer_border = static_cast<MBint>( layerNumber ) - 1;

				bool delete_source_image = false;

				while (true)
				{
					if (current_level < last_layer && src->cols >= 2 && src->rows >= 2)
					{
						dst = new Mat();
						pyrDown( *src, *dst, Size( src->cols / 2, src->rows / 2 ) );

						if (insert_pair.first->second.back() == 0)
						{
							// The last level was not needed -> delete source image
							delete_source_image = true;
						}

						if (current_level == exclusive_border || current_level == layer_border)
						{
							insert_pair.first->second.push_back( dst );
							image_memory_consumption += 3 * dst->cols * dst->rows * sizeof( unsigned char );
							//image_memory_consumption += 3 * dst->cols * dst->rows * sizeof( double );
						}
						else if (current_level < exclusive_border || current_level < layer_border)
						{
							// Save Null-Pointer for not needed Levels
							insert_pair.first->second.push_back( 0 );
						}

						if (delete_source_image)
						{
							delete src;
							delete_source_image = false;
						}

						src = dst;

						current_level++;
					}
					else
					{
						break;
					}
				}

				map<string, MBuint>::iterator im_to_order_it = mImagePathToCurrentOrderMap.find( imagePath );
				map<MBuint, MBImagePathWithMemoryConsumption>::iterator order_to_im_it = mCurrentOrderToImagePathMap.find( im_to_order_it->second );

				im_to_order_it->second = mCurrentImageOrder;
				mCurrentOrderToImagePathMap.insert( make_pair( mCurrentImageOrder++, MBImagePathWithMemoryConsumption( imagePath, order_to_im_it->second.mMemoryConsumption + image_memory_consumption ) ) );
				mCurrentOrderToImagePathMap.erase( order_to_im_it );

				// Check memory consumption, and free some if needed
				mCurrentMemoryConsumption += image_memory_consumption;

				MBuint tries = 0u;
				map<MBuint, MBImagePathWithMemoryConsumption>::iterator order_to_image_it = mCurrentOrderToImagePathMap.begin();

				while (mCurrentMemoryConsumption > mMaxMemoryConsumption && mImagePathToCurrentOrderMap.size() - tries > 1u)
				{
					if (untouchableImagePaths)
					{
						// There are images, that MUST NOT be deleted
						if (untouchableImagePaths->find( order_to_image_it ->second.mImagePath ) == untouchableImagePaths->end() )
						{
							// OK, this image can be deleted
							FreeImage( order_to_image_it->second.mImagePath );
							order_to_image_it = mCurrentOrderToImagePathMap.begin();
						}
						else
						{
							// This image must not be deleted, choose the next one
							order_to_image_it++;
						}
					}
					else
					{
						// No untouchable images, just delete
						FreeImage( mCurrentOrderToImagePathMap.begin()->second.mImagePath );
					}
					tries++;
				}

				if (layerNumber >= insert_pair.first->second.size())
				{
					return 0;
				}
				else
				{
					return insert_pair.first->second[layerNumber];
				}
			}
		}
		else
		{
			// layerNumber < insert_pair.first->second.size()

			// Set the current image to the front of the line

			if (insert_pair.first->second[layerNumber])
			{
				// Existing level
				map<string, MBuint>::iterator im_to_order_it = mImagePathToCurrentOrderMap.find( imagePath );
				map<MBuint, MBImagePathWithMemoryConsumption>::iterator order_to_im_it = mCurrentOrderToImagePathMap.find( im_to_order_it->second );

				im_to_order_it->second = mCurrentImageOrder;
				mCurrentOrderToImagePathMap.insert( make_pair( mCurrentImageOrder++, order_to_im_it->second ) );
				mCurrentOrderToImagePathMap.erase( order_to_im_it );

				return insert_pair.first->second[layerNumber];
			}
			else
			{
				// Non existing level -> free and load again
				FreeImage( imagePath );
				return GetImage( imagePath, layerNumber, untouchableImagePaths );
			}
		}
	}
}

/**
  * Returns the accumulative number of loaded image files
  * Use the ResetNumberOfLoadedImages() method to reset this value to 0
  *
  * @return accumulative number of loaded image files
  */
MBuint MBImageManager::GetNumberOfLoadedImages( void ) const
{
	return mNumImagesLoaded;
}

/**
  * Returns the accumulative time in seconds that was used to load image files
  * Use the ResetTimeNeededForLoading() method to reset this value to 0.0
  *
  * @return accumulative time in seconds that was used to load image files
  */
double MBImageManager::GetTimeInSecondsUsedForImageLoading( void ) const
{
	return mSecondsNeededForLoading;
}

/**
  * Resets the accumulative number of images that were loaded to 0
  */
void MBImageManager::ResetNumberOfLoadedImages( void )
{
	mNumImagesLoaded = 0u;
}

/**
  * Resets the accumulative time in seconds that was used to load image files to 0.0
  */
void MBImageManager::ResetTimeNeededForLoading( void )
{
	mSecondsNeededForLoading = 0.0;
}

/**
  *	Sets the exclusive image pyramid level. After calling this method with a
  * number e.g. 3, the GetImage() method will only store the level 3 and the
  * level that is passed as an argument of a particular image. By setting the
  * level number to the right level, the performance can be increased.
  * The default level is 0.
  *
  * @param exclusiveLevel exclusive image pyramid level
  */
void MBImageManager::SetExclusiveLevel( const MBuint &exclusiveLevel )
{
	mExclusiveLevel = exclusiveLevel;
}

/**
  *	Sets the maximum memory consumption of all images in bytes
  *
  * @param maxMemoryConsumption maximum memory consumption of all images in bytes
  */
void MBImageManager::SetMaximumMemoryConsumptionInBytes( const MBuint &maxMemoryConsumption )
{
	mMaxMemoryConsumption = maxMemoryConsumption;
}